from __future__ import annotations

import asyncio

from neo4j_graphrag.embeddings.openai import OpenAIEmbeddings
from neo4j_graphrag.experimental.components.embedder import TextChunkEmbedder
from my_extractor.openai_ere import (
    LLMEntityRelationExtractor,
    OnError,
)
from neo4j_graphrag.experimental.components.resolver import (
    SinglePropertyExactMatchResolver,
)
from neo4j_graphrag.experimental.components.kg_writer import Neo4jWriter
from neo4j_graphrag.experimental.components.schema import (
    SchemaBuilder,
    NodeType,
    PropertyType,
    RelationshipType,
)
from langchain_text_splitters import RecursiveCharacterTextSplitter
from neo4j_graphrag.experimental.components.text_splitters.langchain import LangChainTextSplitterAdapter

from neo4j_graphrag.experimental.pipeline import Pipeline
from neo4j_graphrag.experimental.pipeline.pipeline import PipelineResult

import neo4j


max_len = 400
overlap = 50
# 文本拆分器
sep_patterns = [
    '```\n', r'\n\*\*\*+\n', r'\n---+\n', r'\n___+\n',
    r'\n\n', r'\n', "。|！|？", r"\.\s|\!\s|\?\s", r"；|;\s"
]


async def define_and_run_pipeline(
        neo4j_driver: neo4j.Driver, openai_llm_config
) -> PipelineResult:
    pipe = Pipeline()
    # define the components
    text_splitter = LangChainTextSplitterAdapter(RecursiveCharacterTextSplitter(
        separators=sep_patterns,
        is_separator_regex=True,
        chunk_size=max_len,
        chunk_overlap=overlap,
        strip_whitespace=True
    ))
    embedder = OpenAIEmbeddings(api_key='EMPTY', base_url='http://172.16.3.123:11434/v1', model='text-embedding')
    resolver = SinglePropertyExactMatchResolver(neo4j_driver)
    pipe.add_component(
        text_splitter,
        "splitter",
    )
    pipe.add_component(TextChunkEmbedder(embedder=embedder), "chunk_embedder")
    pipe.add_component(SchemaBuilder(), "schema")
    pipe.add_component(
        LLMEntityRelationExtractor(
            llm_config=openai_llm_config,
            on_error=OnError.RAISE,
        ),
        "extractor",
    )
    pipe.add_component(Neo4jWriter(neo4j_driver), "writer")
    pipe.add_component(resolver, "resolver")
    # define the execution order of component
    # and how the output of previous components must be used
    pipe.connect("splitter", "chunk_embedder", input_config={"text_chunks": "splitter"})
    pipe.connect("schema", "extractor", input_config={"schema": "schema"})
    pipe.connect(
        "chunk_embedder", "extractor", input_config={"chunks": "chunk_embedder"}
    )
    pipe.connect(
        "extractor",
        "writer",
        input_config={"graph": "extractor"},
    )
    # user input:
    # the initial text
    # and the list of entities and relations we are looking for
    pipe_inputs = {
        "splitter": {
            "text": """阿尔伯特·爱因斯坦是生于1879年的德国物理学家，他撰写了许多开创性的论文，特别是在广义相对论和量子力学领域。他曾在多所机构工作，包括瑞士的伯尔尼大学和英国的牛津大学。"""
        },
        "schema": {
            "node_types": [
                NodeType(
                    label="人",
                    properties=[
                        PropertyType(name="姓名", type="STRING"),
                        PropertyType(name="出生地", type="STRING"),
                        PropertyType(name="生日", type="DATE"),
                    ],
                ),
                NodeType(
                    label="组织",
                    properties=[
                        PropertyType(name="组织名", type="STRING"),
                        PropertyType(name="国家", type="STRING"),
                    ],
                ),
                NodeType(
                    label="区域",
                    properties=[
                        PropertyType(name="区域名", type="STRING"),
                    ],
                ),
            ],
            "relationship_types": [
                RelationshipType(
                    label="工作于",
                ),
                RelationshipType(
                    label="受聘于",
                ),
            ],
            "patterns": [
                ("人", "工作于", "区域"),
                ("人", "受聘于", "组织"),
            ],
        },
        "extractor": {
            "document_info": {
                "path": "my text",
            }
        },
    }
    # run the pipeline
    return await pipe.run(pipe_inputs)


async def main() -> PipelineResult:
    openai_llm_config = {'init': {'base_url': 'http://172.16.3.116:2380/v1', 'api_key': 'EMPTY'},
                         'model': 'Qwen3-32B-AWQ', 'enable_thinking': True}

    driver = neo4j.GraphDatabase.driver(
        "bolt://172.16.3.123:7687", auth=("neo4j", "1234321ly")
    )

    res = await define_and_run_pipeline(driver, openai_llm_config)
    driver.close()
    return res


if __name__ == "__main__":
    res = asyncio.run(main())
    print(res)
